import numpy as np
import scipy.io.wavfile as wf
import python_speech_features as sf
import sklearn.preprocessing as sp
import os
from keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv2D, MaxPooling2D, Flatten, BatchNormalization
from tensorflow.keras.optimizers import Adam

# 整理样本
def search_files(directory):
    files_dict = {}
    for cur_dir, sub_dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".wav"):
                label = cur_dir.split(os.path.sep)[-1]
                if label not in files_dict:
                    files_dict[label] = []
                files_dict[label].append(os.path.join(cur_dir, file))
    return files_dict

def files_mfc(file_urls):
    x_data, y_data = [], []
    for label, urls in file_urls.items():
        for file in urls:
            sample_rate, signs = wf.read(file)
            mfc = sf.mfcc(signs, sample_rate)
            x_data.append(np.mean(mfc, axis=0))
            y_data.append(label)
    return np.array(x_data), y_data

# 读取文件路径
train_urls = search_files("D:/Pycharm/trainvoice5/train")
test_urls = search_files("D:/Pycharm/trainvoice5/test")

# 整理数据
train_x, train_y = files_mfc(train_urls)
test_x, test_y = files_mfc(test_urls)

# 预处理数据
input_shape = (13, 1, 1) # MFCC特征数量为13
train_x = train_x.reshape(train_x.shape[0], *input_shape)
test_x = test_x.reshape(test_x.shape[0], *input_shape)

# 编码标签
encoder = sp.LabelEncoder()
train_y = encoder.fit_transform(train_y)
test_y = encoder.transform(test_y)

# 构建简化的ResNet模型
model = Sequential([
    # 使用 (1, 1) 大小的卷积核
    Conv2D(32, (1, 1), activation='relu', input_shape=input_shape),
    BatchNormalization(),
    MaxPooling2D(pool_size=(2, 1)),
    Flatten(),
    Dense(128, activation='relu'),
    Dense(7, activation='softmax') # 分类数量为7





])
# 编译模型
model.compile(optimizer=Adam(), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(train_x, train_y, epochs=10, batch_size=32, validation_split=0.2)

# 预测测试集
prd_test_y = model.predict(test_x)
prd_test_y = np.argmax(prd_test_y, axis=1)

# 输出分类报告
from sklearn.metrics import classification_report
print(classification_report(test_y, prd_test_y))
